/*
 * Copyright 2023 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package androidx.camera.testing.impl;

import static android.graphics.BitmapFactory.decodeByteArray;
import static android.graphics.ImageFormat.JPEG;
import static android.graphics.ImageFormat.JPEG_R;
import static android.graphics.ImageFormat.RAW_SENSOR;
import static android.graphics.ImageFormat.YUV_420_888;

import static androidx.camera.core.internal.utils.SizeUtil.RESOLUTION_VGA;
import static androidx.camera.testing.impl.ImageProxyUtil.YUV_FORMAT_PLANE_DATA_TYPE_I420;
import static androidx.camera.testing.impl.ImageProxyUtil.createRawImagePlanes;
import static androidx.camera.testing.impl.ImageProxyUtil.createYUV420ImagePlanes;
import static androidx.camera.testing.impl.ImageProxyUtil.getDefaultYuvFormatPlaneDataType;
import static androidx.core.util.Preconditions.checkState;

import android.graphics.Bitmap;
import android.graphics.Canvas;
import android.graphics.Color;
import android.graphics.Gainmap;
import android.graphics.Matrix;
import android.graphics.Paint;
import android.graphics.Rect;
import android.os.Build;

import androidx.annotation.RequiresApi;
import androidx.annotation.VisibleForTesting;
import androidx.camera.core.ImageInfo;
import androidx.camera.core.ImageProxy;
import androidx.camera.core.internal.CameraCaptureResultImageInfo;
import androidx.camera.testing.fakes.FakeCameraCaptureResult;
import androidx.camera.testing.impl.fakes.FakeImageProxy;
import androidx.camera.testing.impl.fakes.FakeJpegPlaneProxy;

import org.jspecify.annotations.NonNull;

import java.io.ByteArrayOutputStream;
import java.nio.ByteBuffer;

/**
 * Generates images for testing.
 *
 * <p> The images generated by this class contains 4 color blocks follows the pattern below. Each
 * block have the same size and covers 1/4 of the image.
 *
 * <pre>
 * ------------------
 * |   red  | green |
 * ------------------
 * |  blue | yellow |
 * ------------------
 * </pre>
 *
 * <p> The gain map generated by this class contains 4 gray scale blocks follows the pattern below.
 * Each block have the same size and covers 1/4 of the image.
 *
 * <pre>
 * ------------------
 * |      black     |
 * ------------------
 * |    dark gray   |
 * ------------------
 * |      gray      |
 * ------------------
 * |      white     |
 * ------------------
 * </pre>
 */
public class TestImageUtil {

    @VisibleForTesting
    public static final int COLOR_BLACK = 0xFF000000;
    @VisibleForTesting
    public static final int COLOR_DARK_GRAY = 0xFF404040;
    @VisibleForTesting
    public static final int COLOR_GRAY = 0xFF808080;
    @VisibleForTesting
    public static final int COLOR_WHITE = 0xFFFFFFFF;

    private TestImageUtil() {
    }

    /**
     * Creates a [FakeImageProxy] with YUV format.
     */
    public static @NonNull FakeImageProxy createYuvFakeImageProxy(@NonNull ImageInfo imageInfo,
            int width, int height) {
        return createYuvFakeImageProxy(imageInfo, width, height, getDefaultYuvFormatPlaneDataType(
                RESOLUTION_VGA.getWidth(), RESOLUTION_VGA.getHeight()), false);
    }

    /**
     * Creates a [FakeImageProxy] with YUV format with the content of the image to match the
     * value of {@link #createBitmap}.
     */
    public static @NonNull FakeImageProxy createYuvFakeImageProxy(@NonNull ImageInfo imageInfo,
            int width, int height,
            @ImageProxyUtil.YuvFormatPlaneDataType int yuvFormatPlaneDataType,
            boolean insertRgbTestData) {
        return createYuvFakeImageProxy(
                createYUV420ImagePlanes(width, height, yuvFormatPlaneDataType, false),
                imageInfo, width, height, 1,
                yuvFormatPlaneDataType == YUV_FORMAT_PLANE_DATA_TYPE_I420 ? 1 : 2,
                insertRgbTestData);
    }

    /**
     * Creates a [FakeImageProxy] with YUV format with the content of the image to match the
     * value of {@link #createBitmap}.
     */
    public static @NonNull FakeImageProxy createYuvFakeImageProxy(@NonNull ImageInfo imageInfo,
            int width, int height, int pixelStrideY, int pixelStrideUV, boolean flipUV,
            boolean insertRgbTestData) {
        return createYuvFakeImageProxy(
                createYUV420ImagePlanes(width, height, pixelStrideY, pixelStrideUV, flipUV, false),
                imageInfo, width, height, pixelStrideY, pixelStrideUV, insertRgbTestData);
    }

    private static @NonNull FakeImageProxy createYuvFakeImageProxy(
            ImageProxy.@NonNull PlaneProxy[] planeProxies, @NonNull ImageInfo imageInfo, int width,
            int height, int pixelStrideY, int pixelStrideUV, boolean insertRgbTestData) {
        FakeImageProxy image = new FakeImageProxy(imageInfo);
        image.setFormat(YUV_420_888);
        image.setPlanes(planeProxies);
        image.setWidth(width);
        image.setHeight(height);

        // Directly returns the image if RGB test data is not needed.
        if (!insertRgbTestData) {
            return image;
        }

        Bitmap rgbBitmap = createBitmap(width, height);
        writeBitmapToYuvByteBuffers(rgbBitmap,
                image.getPlanes()[0].getBuffer(),
                image.getPlanes()[1].getBuffer(),
                image.getPlanes()[2].getBuffer(),
                pixelStrideY, pixelStrideUV);

        return image;
    }

    private static void writeBitmapToYuvByteBuffers(
            @NonNull Bitmap bitmap,
            @NonNull ByteBuffer yByteBuffer,
            @NonNull ByteBuffer uByteBuffer,
            @NonNull ByteBuffer vByteBuffer,
            int pixelStrideY,
            int pixelStrideUV) {
        int width = bitmap.getWidth();
        int height = bitmap.getHeight();
        int[] argb = new int[width * height];
        bitmap.getPixels(argb, 0, width, 0, 0, width, height);

        int yIndex = 0;
        int uvIndex = 0;

        for (int j = 0; j < height; j++) {
            for (int i = 0; i < width; i++) {
                int rgb = argb[j * width + i];
                int r = (rgb >> 16) & 0xFF;
                int g = (rgb >> 8) & 0xFF;
                int b = rgb & 0xFF;

                int y = (int) (0.299 * r + 0.587 * g + 0.114 * b);
                int u = (int) (-0.169 * r - 0.331 * g + 0.5 * b + 128);
                int v = (int) (0.5 * r - 0.419 * g - 0.081 * b + 128);

                yByteBuffer.put(yIndex, (byte) y);
                yIndex += pixelStrideY;
                if (j % 2 == 0 && i % 2 == 0) {
                    uByteBuffer.put(uvIndex * pixelStrideUV, (byte) u);
                    vByteBuffer.put(uvIndex * pixelStrideUV, (byte) v);
                    uvIndex++;
                }
            }
        }
        yByteBuffer.rewind();
        uByteBuffer.rewind();
        vByteBuffer.rewind();
    }

    /**
     * Creates a [FakeImageProxy] with [RAW_SENSOR] format.
     */
    public static @NonNull FakeImageProxy createRawFakeImageProxy(@NonNull ImageInfo imageInfo,
            int width, int height) {
        FakeImageProxy image = new FakeImageProxy(imageInfo);
        image.setFormat(RAW_SENSOR);
        image.setPlanes(createRawImagePlanes(width, height, 2, false));
        image.setWidth(width);
        image.setHeight(height);
        return image;
    }

    /**
     * Creates a {@link FakeImageProxy} from JPEG bytes.
     */
    public static @NonNull FakeImageProxy createJpegFakeImageProxy(@NonNull ImageInfo imageInfo,
            byte @NonNull [] jpegBytes) {
        Bitmap bitmap = decodeByteArray(jpegBytes, 0, jpegBytes.length);
        return createJpegFakeImageProxy(imageInfo, jpegBytes, bitmap.getWidth(),
                bitmap.getHeight());
    }

    /**
     * Creates a {@link FakeImageProxy} from JPEG bytes of JPEG/R.
     */
    public static @NonNull FakeImageProxy createJpegrFakeImageProxy(@NonNull ImageInfo imageInfo,
            byte @NonNull [] jpegBytes) {
        Bitmap bitmap = decodeByteArray(jpegBytes, 0, jpegBytes.length);
        return createJpegrFakeImageProxy(imageInfo, jpegBytes, bitmap.getWidth(),
                bitmap.getHeight());
    }

    /**
     * Creates a {@link FakeImageProxy} from JPEG bytes.
     */
    public static @NonNull FakeImageProxy createJpegFakeImageProxy(@NonNull ImageInfo imageInfo,
            byte @NonNull [] jpegBytes, int width, int height) {
        FakeImageProxy image = new FakeImageProxy(imageInfo);
        image.setFormat(JPEG);
        image.setPlanes(new FakeJpegPlaneProxy[]{new FakeJpegPlaneProxy(jpegBytes)});
        image.setWidth(width);
        image.setHeight(height);
        return image;
    }

    /**
     * Creates a {@link FakeImageProxy} from JPEG bytes of JPEG/R.
     */
    public static @NonNull FakeImageProxy createJpegrFakeImageProxy(@NonNull ImageInfo imageInfo,
            byte @NonNull [] jpegBytes, int width, int height) {
        FakeImageProxy image = new FakeImageProxy(imageInfo);
        image.setFormat(JPEG_R);
        image.setPlanes(new FakeJpegPlaneProxy[]{new FakeJpegPlaneProxy(jpegBytes)});
        image.setWidth(width);
        image.setHeight(height);
        return image;
    }

    /**
     * Creates a {@link FakeImageProxy} from JPEG bytes.
     */
    public static @NonNull FakeImageProxy createJpegFakeImageProxy(byte @NonNull [] jpegBytes) {
        return createJpegFakeImageProxy(
                new CameraCaptureResultImageInfo(new FakeCameraCaptureResult()), jpegBytes);
    }

    /**
     * Creates a {@link FakeImageProxy} from JPEG bytes of JPEG/R.
     */
    public static @NonNull FakeImageProxy createJpegrFakeImageProxy(byte @NonNull [] jpegBytes) {
        return createJpegrFakeImageProxy(
                new CameraCaptureResultImageInfo(new FakeCameraCaptureResult()), jpegBytes);
    }

    /**
     * Generates a JPEG image.
     */
    public static byte @NonNull [] createJpegBytes(int width, int height) {
        Bitmap bitmap = createBitmap(width, height);
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        bitmap.compress(Bitmap.CompressFormat.JPEG, 100, outputStream);
        return outputStream.toByteArray();
    }

    /**
     * Generates a JPEG/R image.
     */
    @RequiresApi(api = Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
    public static byte @NonNull [] createJpegrBytes(int width, int height) {
        Bitmap bitmap = createBitmapWithGainmap(width, height);
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        bitmap.compress(Bitmap.CompressFormat.JPEG, 100, outputStream);
        return outputStream.toByteArray();
    }

    /**
     * Generates a A24 problematic JPEG image.
     */
    public static byte @NonNull [] createA24ProblematicJpegByteArray(int width, int height) {
        byte[] incorrectHeaderByteData =
                new byte[]{(byte) 0xff, (byte) 0xd8, (byte) 0xff, (byte) 0xe1, (byte) 0xff,
                        (byte) 0x7c, (byte) 0x45, (byte) 0x78, (byte) 0x69, (byte) 0x66,
                        (byte) 0x00, (byte) 0x00};
        byte[] jpegBytes = createJpegBytes(width, height);
        byte[] result = new byte[incorrectHeaderByteData.length + jpegBytes.length];
        System.arraycopy(incorrectHeaderByteData, 0, result, 0, incorrectHeaderByteData.length);
        System.arraycopy(jpegBytes, 0, result, incorrectHeaderByteData.length, jpegBytes.length);
        return result;
    }

    /**
     * Generates a {@link Bitmap} image and paints it with 4 color blocks.
     */
    public static @NonNull Bitmap createBitmap(int width, int height) {
        Bitmap bitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888);
        int centerX = width / 2;
        int centerY = height / 2;
        Canvas canvas = new Canvas(bitmap);
        canvas.drawRect(0, 0, centerX, centerY, createPaint(Color.RED));
        canvas.drawRect(centerX, 0, width, centerY, createPaint(Color.GREEN));
        canvas.drawRect(centerX, centerY, width, height, createPaint(Color.YELLOW));
        canvas.drawRect(0, centerY, centerX, height, createPaint(Color.BLUE));
        return bitmap;
    }

    /**
     * Generates a {@link Bitmap} image (contains 4 color blocks) with gain map (contains 4 gray
     * scale blocks) set.
     */
    @RequiresApi(api = Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
    public static @NonNull Bitmap createBitmapWithGainmap(int width, int height) {
        Bitmap bitmap = createBitmap(width, height);
        Api34Impl.setGainmap(bitmap, createGainmap(width, height));
        return bitmap;
    }

    /**
     * Generates a {@link Gainmap} image and paints it with 4 gray scale blocks.
     */
    @RequiresApi(api = Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
    public static @NonNull Gainmap createGainmap(int width, int height) {
        Bitmap bitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888);
        int oneFourthY = height / 4;
        int twoFourthsY = oneFourthY * 2;
        int threeFourthsY = oneFourthY * 3;
        Canvas canvas = new Canvas(bitmap);
        canvas.drawRect(0, 0, width, oneFourthY, createPaint(COLOR_BLACK));
        canvas.drawRect(0, oneFourthY, width, twoFourthsY, createPaint(COLOR_DARK_GRAY));
        canvas.drawRect(0, twoFourthsY, width, threeFourthsY, createPaint(COLOR_GRAY));
        canvas.drawRect(0, threeFourthsY, width, height, createPaint(COLOR_WHITE));
        return Api34Impl.createGainmap(bitmap);
    }

    /**
     * Rotates the bitmap clockwise by the given degrees.
     */
    public static @NonNull Bitmap rotateBitmap(@NonNull Bitmap bitmap, int rotationDegrees) {
        Matrix matrix = new Matrix();
        matrix.postRotate(rotationDegrees);
        return Bitmap.createBitmap(bitmap, 0, 0, bitmap.getWidth(), bitmap.getHeight(), matrix,
                true);
    }

    /**
     * Calculates the average color difference between the 2 JPEG images.
     */
    public static int getAverageDiff(byte @NonNull [] jpeg1, byte @NonNull [] jpeg2) {
        return getAverageDiff(
                decodeByteArray(jpeg1, 0, jpeg1.length),
                decodeByteArray(jpeg2, 0, jpeg2.length));
    }

    /**
     * Calculates the average color difference between the 2 bitmaps.
     */
    public static int getAverageDiff(@NonNull Bitmap bitmap1, @NonNull Bitmap bitmap2) {
        checkState(bitmap1.getWidth() == bitmap2.getWidth());
        checkState(bitmap1.getHeight() == bitmap2.getHeight());
        int totalDiff = 0;
        for (int i = 0; i < bitmap1.getWidth(); i++) {
            for (int j = 0; j < bitmap1.getHeight(); j++) {
                totalDiff += calculateColorDiff(bitmap1.getPixel(i, j), bitmap2.getPixel(i, j));
            }
        }
        return totalDiff / (bitmap1.getWidth() * bitmap2.getHeight());
    }

    /**
     * Calculates the average color difference, between the given image/crop rect and the color.
     *
     * <p>This method is used for checking the content of an image is correct.
     */
    public static int getAverageDiff(@NonNull Bitmap bitmap, @NonNull Rect rect, int color) {
        int totalDiff = 0;
        for (int i = rect.left; i < rect.right; i++) {
            for (int j = rect.top; j < rect.bottom; j++) {
                totalDiff += calculateColorDiff(bitmap.getPixel(i, j), color);
            }
        }
        return totalDiff / (rect.width() * rect.height());
    }

    /**
     * Calculates the difference between 2 colors.
     *
     * <p>The difference is calculated as the average difference of each R, G and B color
     * components.
     */
    private static int calculateColorDiff(int color1, int color2) {
        int diff = 0;
        for (int shift = 0; shift <= 16; shift += 8) {
            diff += Math.abs(((color1 >> shift) & 0xFF) - ((color2 >> shift) & 0xFF));
        }
        return diff / 3;
    }

    /**
     * Creates a FILL paint with the given color.
     */
    private static Paint createPaint(int color) {
        Paint paint = new Paint();
        paint.setStyle(Paint.Style.FILL);
        paint.setColor(color);
        return paint;
    }

    @RequiresApi(34)
    private static class Api34Impl {
        static Gainmap createGainmap(@NonNull Bitmap bitmap) {
            return new Gainmap(bitmap);
        }

        static void setGainmap(@NonNull Bitmap bitmap, @NonNull Gainmap gainmap) {
            bitmap.setGainmap(gainmap);
        }

        // This class is not instantiable.
        private Api34Impl() {
        }
    }
}
